import time

import requests
import base64
import ddddocr
from lxml import html
from src import utils

ocr = ddddocr.DdddOcr()

def adapter(imgCode):
    if 'H' in imgCode:
        return imgCode.replace('H', '11')
    if 'o' in imgCode:
        return imgCode.replace('o', '0')
    if 'O' in imgCode:
        return imgCode.replace('O', '0')
    if 'l' in imgCode:
        return imgCode.replace('l', '1')

    return imgCode

if __name__ == '__main__':
    url = 'https://spain.blscn.cn/CHN/CaptchaPublic/GenerateCaptcha?data=tE%2f7gWSSTw4v%2baL4MBy6jU8dTrmds3sAgQqRQqX5WZ5LWML6WbSb4TBSbI%2btwSUBVT8G5pQpWWapa6PI2HzKQHSN%2ftya03U%2bXR%2fhy0bn6%2bE%3d'
    htmlResponse = requests.get(url, headers=utils.get_headers())
    if htmlResponse.status_code == 200:
        # 将二进制内容写入本地文件
        tree = html.fromstring(htmlResponse.content)
        images = tree.xpath('//img[@class="captcha-img"]')
        for img in images:
            img_src = img.get('src')
            img_data = base64.b64decode(img_src.replace('data:image/gif;base64,', ''))
            imgCode = ocr.classification(img_data) or 'none'
            file_name = adapter(imgCode) + '_' + str(time.time()) + '.png'
            with open('train/' + file_name, 'wb') as file:
                file.write(img_data)
